import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor()
])

# 加载数据集
dataset_train = ImageFolder(root='F:/深度学习数据集/imagenet-10/train/', transform=transform)
dataset_val = ImageFolder(root='F:/深度学习数据集/imagenet-10/test/', transform=transform)
data_loader_train = DataLoader(dataset_train, batch_size=64, shuffle=False, num_workers=8)
data_loader_val = DataLoader(dataset_val, batch_size=64, shuffle=False, num_workers=8)

if __name__ == "__main__":
    # 计算均值和标准差
    mean = 0.
    std = 0.
    nb_samples = 0.
    for data, _ in data_loader_train:
        batch_samples = data.size(0)
        data = data.view(batch_samples, data.size(1), -1)
        mean += data.mean(2).sum(0)
        std += data.std(2).sum(0)
        nb_samples += batch_samples
    for data, _ in data_loader_val:
        batch_samples = data.size(0)
        data = data.view(batch_samples, data.size(1), -1)
        mean += data.mean(2).sum(0)
        std += data.std(2).sum(0)
        nb_samples += batch_samples
    mean /= nb_samples
    std /= nb_samples

    print(f'Mean: {mean}')
    print(f'Std: {std}')